{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install chromadb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 初始化智谱 embedding\n",
    "from typing import Any, Dict, List, Optional\n",
    "\n",
    "from langchain_core.embeddings import Embeddings\n",
    "from langchain_core.utils import get_from_dict_or_env\n",
    "from pydantic import BaseModel, Field, model_validator\n",
    "import os\n",
    "from zhipuai import ZhipuAI\n",
    "\n",
    "class ZhipuAIEmbeddings(BaseModel, Embeddings):\n",
    "    \"\"\"ZhipuAI embedding model integration.\n",
    "\n",
    "    Setup:\n",
    "\n",
    "        To use, you should have the ``zhipuai`` python package installed, and the\n",
    "        environment variable ``ZHIPU_API_KEY`` set with your API KEY.\n",
    "\n",
    "        More instructions about ZhipuAi Embeddings, you can get it\n",
    "        from  https://open.bigmodel.cn/dev/api#vector\n",
    "\n",
    "        .. code-block:: bash\n",
    "\n",
    "            pip install -U zhipuai\n",
    "            export ZHIPU_API_KEY=\"your-api-key\"\n",
    "\n",
    "    Key init args — completion params:\n",
    "        model: Optional[str]\n",
    "            Name of ZhipuAI model to use.\n",
    "        api_key: str\n",
    "            Automatically inferred from env var `ZHIPU_API_KEY` if not provided.\n",
    "\n",
    "    See full list of supported init args and their descriptions in the params section.\n",
    "\n",
    "    Instantiate:\n",
    "\n",
    "        .. code-block:: python\n",
    "\n",
    "            from langchain_community.embeddings import ZhipuAIEmbeddings\n",
    "\n",
    "            embed = ZhipuAIEmbeddings(\n",
    "                model=\"embedding-2\",\n",
    "                # api_key=\"...\",\n",
    "            )\n",
    "\n",
    "    Embed single text:\n",
    "        .. code-block:: python\n",
    "\n",
    "            input_text = \"The meaning of life is 42\"\n",
    "            embed.embed_query(input_text)\n",
    "\n",
    "        .. code-block:: python\n",
    "\n",
    "            [-0.003832892, 0.049372625, -0.035413884, -0.019301128, 0.0068899863, 0.01248398, -0.022153955, 0.006623926, 0.00778216, 0.009558191, ...]\n",
    "\n",
    "\n",
    "    Embed multiple text:\n",
    "        .. code-block:: python\n",
    "\n",
    "            input_texts = [\"This is a test query1.\", \"This is a test query2.\"]\n",
    "            embed.embed_documents(input_texts)\n",
    "\n",
    "        .. code-block:: python\n",
    "\n",
    "            [\n",
    "                [0.0083934665, 0.037985895, -0.06684559, -0.039616987, 0.015481004, -0.023952313, ...],\n",
    "                [-0.02713102, -0.005470169, 0.032321047, 0.042484466, 0.023290444, 0.02170547, ...]\n",
    "            ]\n",
    "    \"\"\"  # noqa: E501\n",
    "\n",
    "    client: Any = Field(default=None, exclude=True)  #: :meta private:\n",
    "    model: str = Field(default=\"embedding-2\")\n",
    "    \"\"\"Model name\"\"\"\n",
    "    api_key: str\n",
    "    \"\"\"Automatically inferred from env var `ZHIPU_API_KEY` if not provided.\"\"\"\n",
    "    dimensions: Optional[int] = None\n",
    "    \"\"\"The number of dimensions the resulting output embeddings should have.\n",
    "\n",
    "    Only supported in `embedding-3` and later models.\n",
    "    \"\"\"\n",
    "\n",
    "    @model_validator(mode=\"before\")\n",
    "    @classmethod\n",
    "    def validate_environment(cls, values: Dict) -> Any:\n",
    "        \"\"\"Validate that auth token exists in environment.\"\"\"\n",
    "        values[\"api_key\"] = get_from_dict_or_env(values, \"api_key\", \"ZHIPUAI_API_KEY\")\n",
    "        try:\n",
    "            from zhipuai import ZhipuAI\n",
    "\n",
    "            values[\"client\"] = ZhipuAI(api_key=values[\"api_key\"])\n",
    "        except ImportError:\n",
    "            raise ImportError(\n",
    "                \"Could not import zhipuai python package.\"\n",
    "                \"Please install it with `pip install zhipuai`.\"\n",
    "            )\n",
    "        return values\n",
    "\n",
    "\n",
    "\n",
    "    def embed_query(self, text: str) -> List[float]:\n",
    "        \"\"\"\n",
    "        Embeds a text using the AutoVOT algorithm.\n",
    "\n",
    "        Args:\n",
    "            text: A text to embed.\n",
    "\n",
    "        Returns:\n",
    "            Input document's embedded list.\n",
    "        \"\"\"\n",
    "        resp = self.embed_documents([text])\n",
    "        return resp[0]\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "    def embed_documents(self, texts: List[str]) -> List[List[float]]:\n",
    "        \"\"\"\n",
    "        Embeds a list of text documents using the AutoVOT algorithm.\n",
    "\n",
    "        Args:\n",
    "            texts: A list of text documents to embed.\n",
    "\n",
    "        Returns:\n",
    "            A list of embeddings for each document in the input list.\n",
    "            Each embedding is represented as a list of float values.\n",
    "        \"\"\"\n",
    "        if self.dimensions is not None:\n",
    "            resp = self.client.embeddings.create(\n",
    "                model=self.model,\n",
    "                input=texts,\n",
    "                dimensions=self.dimensions,\n",
    "            )\n",
    "        else:\n",
    "            resp = self.client.embeddings.create(model=self.model, input=texts)\n",
    "        embeddings = [r.embedding for r in resp.data]\n",
    "        return embeddings\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "import chromadb\n",
    "from chromadb.utils import embedding_functions\n",
    "from pprint import pprint\n",
    "import os\n",
    "\n",
    "# 初始化Chroma客户端（持久化模式）\n",
    "client = chromadb.PersistentClient(path=\"ch16_db\")\n",
    "\n",
    "# 创建一个集合（类似于表）\n",
    "collection = client.get_or_create_collection(\"products\")\n",
    "\n",
    "# 定义嵌入函数（使用预训练的嵌入模型）\n",
    "embedding = ZhipuAIEmbeddings(\n",
    "    model=\"embedding-2\",\n",
    "    api_key=os.getenv(\"ZHIPU_API_KEY\"),\n",
    "    dimensions=1024\n",
    ")\n",
    "\n",
    "# 添加嵌入数据\n",
    "collection.add(\n",
    "    documents=[\"Galaxy S21\", \"iPhone 13\", \"MacBook Pro\"],\n",
    "    embeddings=embedding.embed_documents([\"Galaxy S21\", \"iPhone 13\", \"MacBook Pro\"]),\n",
    "    metadatas=[\n",
    "        {\"category\": \"手机\", \"price\": 799.99},\n",
    "        {\"category\": \"手机\", \"price\": 999.99},\n",
    "        {\"category\": \"笔记本电脑\", \"price\": 1299.99}\n",
    "    ],\n",
    "    ids=[\"prod1\", \"prod2\", \"prod3\"]\n",
    ")\n",
    "\n",
    "print(\"数据添加完成！\")\n",
    "\n",
    "# 获取集合中的所有数据\n",
    "all_data = collection.get()\n",
    "print(\"集合中的所有数据：\")\n",
    "pprint(all_data)\n",
    "\n",
    "# 根据 ID 获取特定的文档\n",
    "specific_data = collection.get(ids=[\"prod1\"])\n",
    "print(\"\\nID 为 'prod1' 的文档：\")\n",
    "pprint(specific_data)\n",
    "\n",
    "# 根据元数据条件获取文档\n",
    "filtered_data = collection.get(where={\"category\": \"手机\"})\n",
    "print(\"\\n类别为 '手机' 的文档：\")\n",
    "pprint(filtered_data)\n",
    "\n",
    "# 更新已有文档的元数据\n",
    "collection.update(\n",
    "    ids=[\"prod1\"],\n",
    "    metadatas=[{\"category\": \"手机\", \"price\": 749.99}]\n",
    ")\n",
    "print(\"\\n已更新 ID 为 'prod1' 的文档价格。\")\n",
    "\n",
    "# 删除特定 ID 的文档\n",
    "collection.delete(ids=[\"prod2\"])\n",
    "print(\"\\n已删除 ID 为 'prod2' 的文档。\")\n",
    "\n",
    "# 查看集合中剩余的文档\n",
    "remaining_data = collection.get()\n",
    "print(\"\\n剩余的文档：\")\n",
    "pprint(remaining_data)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import chromadb\n",
    "from chromadb.utils import embedding_functions\n",
    "from pprint import pprint\n",
    "import os\n",
    "\n",
    "# 初始化Chroma客户端（持久化模式）\n",
    "client = chromadb.PersistentClient(path=\"ch16_db\")\n",
    "\n",
    "# 创建一个集合（类似于表）\n",
    "collection = client.get_or_create_collection(\"products\")\n",
    "\n",
    "# 定义嵌入函数（使用预训练的嵌入模型）\n",
    "embedding = ZhipuAIEmbeddings(\n",
    "    model=\"embedding-2\",\n",
    "    api_key=os.getenv(\"ZHIPU_API_KEY\"),\n",
    "    dimensions=1024\n",
    ")\n",
    "\n",
    "# 添加嵌入数据\n",
    "collection.add(\n",
    "    documents=[\"Galaxy S21\", \"iPhone 13\", \"MacBook Pro\"],\n",
    "    embeddings=embedding.embed_documents([\"Galaxy S21\", \"iPhone 13\", \"MacBook Pro\"]),\n",
    "    metadatas=[\n",
    "        {\"category\": \"手机\", \"price\": 799.99},\n",
    "        {\"category\": \"手机\", \"price\": 999.99},\n",
    "        {\"category\": \"笔记本电脑\", \"price\": 1299.99}\n",
    "    ],\n",
    "    ids=[\"prod1\", \"prod2\", \"prod3\"]\n",
    ")\n",
    "\n",
    "print(\"数据添加完成！\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 获取集合中的所有数据\n",
    "all_data = collection.get()\n",
    "print(\"集合中的所有数据：\")\n",
    "pprint(all_data)\n",
    "\n",
    "# 根据 ID 获取特定的文档\n",
    "specific_data = collection.get(ids=[\"prod1\"])\n",
    "print(\"\\nID 为 'prod1' 的文档：\")\n",
    "pprint(specific_data)\n",
    "\n",
    "# 根据元数据条件获取文档\n",
    "filtered_data = collection.get(where={\"category\": \"手机\"})\n",
    "print(\"\\n类别为 '手机' 的文档：\")\n",
    "pprint(filtered_data)\n",
    "\n",
    "# 更新已有文档的元数据\n",
    "collection.update(\n",
    "    ids=[\"prod1\"],\n",
    "    metadatas=[{\"category\": \"手机\", \"price\": 749.99}]\n",
    ")\n",
    "print(\"\\n已更新 ID 为 'prod1' 的文档价格。\")\n",
    "\n",
    "# 删除特定 ID 的文档\n",
    "collection.delete(ids=[\"prod2\"])\n",
    "print(\"\\n已删除 ID 为 'prod2' 的文档。\")\n",
    "\n",
    "# 查看集合中剩余的文档\n",
    "remaining_data = collection.get()\n",
    "print(\"\\n剩余的文档：\")\n",
    "pprint(remaining_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 导入所需的模块\n",
    "from langchain.vectorstores import Chroma\n",
    "import os\n",
    "\n",
    "# 从磁盘加载持久化数据库\n",
    "persist_directory = \"ch16_db\"\n",
    "embedding = ZhipuAIEmbeddings(\n",
    "    model=\"embedding-2\",\n",
    "    api_key=os.getenv(\"ZHIPU_API_KEY\"),\n",
    "    dimensions=1024\n",
    ")\n",
    "\n",
    "vectordb = Chroma(persist_directory=persist_directory, embedding_function=embedding)\n",
    "\n",
    "# 执行相似性搜索并返回结果及其得分\n",
    "query = \"智能手机\"\n",
    "results = vectordb.similarity_search_with_score(query)\n",
    "\n",
    "# 输出搜索结果\n",
    "print(\"Search results:\")\n",
    "for doc, score in results:\n",
    "    print(f\"Document: {doc.page_content}, Score: {score}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# 创建基于Chroma向量存储的检索器\n",
    "retriever = vectordb.as_retriever(search_type=\"mmr\")\n",
    "\n",
    "# 检索与查询相关的文档\n",
    "retrieved_docs = retriever.get_relevant_documents(query)\n",
    "\n",
    "# 输出第一篇检索到的文档内容\n",
    "print(\"Retrieved document:\")\n",
    "print(retrieved_docs[0].page_content)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 导入所需的模块\n",
    "from langchain.vectorstores import Chroma\n",
    "import os\n",
    "\n",
    "# 从磁盘加载持久化数据库\n",
    "persist_directory = \"ch16_db\"\n",
    "embedding = ZhipuAIEmbeddings(\n",
    "    model=\"embedding-2\",\n",
    "    api_key=os.getenv(\"ZHIPU_API_KEY\"),  # 确保环境变量已设置或直接替换为您的API密钥\n",
    "    dimensions=1024\n",
    ")\n",
    "\n",
    "vectordb = Chroma(persist_directory=persist_directory, embedding_function=embedding)\n",
    "\n",
    "# 执行相似性搜索并返回结果及其得分\n",
    "query = \"智能手机\"\n",
    "result = vectordb.similarity_search_with_score(query)\n",
    "\n",
    "# 输出第一条搜索结果\n",
    "print(result)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 创建基于Chroma向量存储的检索器\n",
    "retriever = vectordb.as_retriever(search_type=\"mmr\")\n",
    "\n",
    "# 检索与查询相关的文档\n",
    "retrieved_docs = retriever.get_relevant_documents(query)\n",
    "\n",
    "# 输出第一篇检索到的文档内容\n",
    "print(retrieved_docs[0].page_content)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "collection = client.get_or_create_collection(\"products\")\n",
    "\n",
    "# 定义嵌入函数（假设使用预训练的嵌入模型）\n",
    "\n",
    "# 添加嵌入数据\n",
    "collection.add(\n",
    "    documents=[\"Galaxy S21\", \"iPhone 13\", \"MacBook Pro\"],\n",
    "    embeddings=embedding.embed_documents([\"Galaxy S21\", \"iPhone 13\", \"MacBook Pro\"]),\n",
    "    metadatas=[\n",
    "        {\"category\": \"手机\", \"price\": 799.99},\n",
    "        {\"category\": \"手机\", \"price\": 999.99},\n",
    "        {\"category\": \"笔记本电脑\", \"price\": 1299.99}\n",
    "    ],\n",
    "    ids=[\"prod1\", \"prod2\", \"prod3\"]\n",
    ")\n",
    "\n",
    "# 查询相似向量\n",
    "\n",
    "\n",
    "\n",
    "results = collection.query(\n",
    "    query_embeddings=embedding.embed_documents([\"智能手机\"]),\n",
    "    n_results=2\n",
    ")\n",
    "\n",
    "pprint(results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(collection.count())  # 确保集合中有文档\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "query_embeddings = embedding.embed_documents([\"智能手机\"])\n",
    "print(query_embeddings)  # 检查查询嵌入是否正确生成"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "embeddings = embedding.embed_documents([\"Galaxy S21\", \"iPhone 13\", \"MacBook Pro\"])\n",
    "print(embeddings)  # 确保嵌入生成成功\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 导入所需的模块\n",
    "from langchain.vectorstores import Chroma\n",
    "\n",
    "# 从磁盘加载持久化数据库\n",
    "persist_directory = \"ch16_db\"\n",
    "embedding = ZhipuAIEmbeddings(\n",
    "    model=\"embedding-2\",\n",
    "    api_key=os.getenv(\"ZHIPU_API_KEY\"),\n",
    "    dimensions=1024\n",
    ")\n",
    "vectordb = Chroma(persist_directory=persist_directory, embedding_function=embedding)\n",
    "# 执行相似性搜索并返回结果及其得分\n",
    "query = \"Galaxy S21\"\n",
    "print(query)\n",
    "result = vectordb.similarity_search_with_score(query)\n",
    "# 输出第一条搜索结果\n",
    "print(result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 导入所需模块\n",
    "from langchain.vectorstores import Chroma\n",
    "from langchain_core.embeddings import Embeddings\n",
    "from zhipuai import ZhipuAI\n",
    "import os\n",
    "\n",
    "\n",
    "# 设置持久化目录\n",
    "persist_directory = \"ch16_db\"\n",
    "\n",
    "\n",
    "\n",
    "# 初始化 Chroma 数据库\n",
    "vectordb = Chroma(persist_directory=persist_directory, embedding_function=embedding)\n",
    "\n",
    "# 执行相似性搜索并返回结果及其得分\n",
    "query = \"iphone\"\n",
    "print(\"Query:\", query)\n",
    "\n",
    "# 生成查询嵌入\n",
    "query_embeddings = embedding.embed_documents([query])\n",
    "print(\"Query embeddings:\", query_embeddings)  # 打印查询嵌入\n",
    "\n",
    "# 获取集合并检查数据库中的文档数量\n",
    "\n",
    "documents = collection.get()  # 获取所有文档\n",
    "document_count = len(documents)  # 计算文档数量\n",
    "print(\"Number of documents in database:\", document_count)\n",
    "\n",
    "# 如果数据库中有文档，则进行相似性搜索\n",
    "if document_count > 0:\n",
    "    # 进行相似性搜索\n",
    "    result = vectordb.similarity_search_with_score(query, k=5)\n",
    "    # 输出结果\n",
    "    print(\"Search results:\", result)\n",
    "else:\n",
    "    print(\"数据库中没有文档，无法进行搜索。\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import chromadb\n",
    "from langchain.vectorstores import Chroma\n",
    "\n",
    "import os\n",
    "\n",
    "# 初始化Chroma客户端\n",
    "persist_directory = \"ch16_db\"\n",
    "embedding = ZhipuAIEmbeddings(model=\"embedding-2\", api_key=os.getenv(\"ZHIPU_API_KEY\"), dimensions=1024)\n",
    "vectordb = Chroma(persist_directory=persist_directory, embedding_function=embedding)\n",
    "\n",
    "# 查询内容\n",
    "query = \"Galaxy S21\"\n",
    "query_embeddings = embedding.embed_documents([query])\n",
    "\n",
    "# 检查查询嵌入\n",
    "print(\"Query embeddings:\", query_embeddings)\n",
    "\n",
    "# 获取集合中的文档\n",
    "documents = vectordb.get()  # 确保从vectordb获取文档\n",
    "document_count = len(documents)\n",
    "print(\"Number of documents in database:\", document_count)\n",
    "\n",
    "# 如果有文档，则进行相似性搜索\n",
    "if document_count > 0:\n",
    "    result = vectordb.similarity_search_with_score(query, k=5)\n",
    "    print(\"Search results:\", result)\n",
    "else:\n",
    "    print(\"数据库中没有文档，无法进行搜索。\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import chromadb\n",
    "from chromadb.utils import embedding_functions\n",
    "from pprint import pprint\n",
    "import os\n",
    "import numpy as np\n",
    "\n",
    "# 假设您已经定义了 ZhipuAIEmbeddings，如果没有，需要替换为实际的嵌入函数\n",
    "# 请确保您已经正确导入或定义了 ZhipuAIEmbeddings 类\n",
    "\n",
    "# 初始化 Chroma 客户端（持久化模式）\n",
    "client = chromadb.PersistentClient(path=\"ch16_db\")\n",
    "\n",
    "# 删除并重新创建集合\n",
    "try:\n",
    "    client.delete_collection(name=\"products\")\n",
    "    print(\"集合已删除。\")\n",
    "except Exception as e:\n",
    "    print(\"集合不存在，创建新的集合。\")\n",
    "\n",
    "# 创建新的集合\n",
    "collection = client.create_collection(name=\"products\")\n",
    "print(\"新的集合已创建。\")\n",
    "\n",
    "# 定义嵌入函数（使用预训练的嵌入模型）\n",
    "embedding = ZhipuAIEmbeddings(\n",
    "    model=\"embedding-2\",\n",
    "    api_key=os.getenv(\"ZHIPU_API_KEY\"),\n",
    "    dimensions=1024\n",
    ")\n",
    "\n",
    "# 添加嵌入数据\n",
    "collection.add(\n",
    "    documents=[\"Galaxy S21\", \"iPhone 13\", \"MacBook Pro\"],\n",
    "    embeddings=embedding.embed_documents([\"Galaxy S21\", \"iPhone 13\", \"MacBook Pro\"]),\n",
    "    metadatas=[\n",
    "        {\"category\": \"手机\", \"price\": 799.99},\n",
    "        {\"category\": \"手机\", \"price\": 999.99},\n",
    "        {\"category\": \"笔记本电脑\", \"price\": 1299.99}\n",
    "    ],\n",
    "    ids=[\"prod1\", \"prod2\", \"prod3\"]\n",
    ")\n",
    "print(\"数据添加完成！\")\n",
    "\n",
    "# 打印集合中的所有数据，验证添加是否成功\n",
    "# all_data = collection.get(include=[\"documents\", \"metadatas\", \"ids\"])  # 错误的方式\n",
    "all_data = collection.get(include=[\"documents\", \"metadatas\"])  # 移除 \"ids\"注意，ids 会自动返回。不需要写到这里面\n",
    "\n",
    "print(\"当前集合中的数据：\")\n",
    "for metadata, document, doc_id in zip(all_data.get('metadatas', []), all_data.get('documents', []), all_data.get('ids', [])):\n",
    "    print(f\"ID: {doc_id}, 产品: {document}, 分类: {metadata['category']}, 价格: {metadata['price']}\")\n",
    "\n",
    "# 自定义元数据过滤：获取价格低于1000的手机，并包括嵌入向量\n",
    "filtered_data = collection.get(\n",
    "    where={\n",
    "        \"$and\": [\n",
    "            {\"category\": \"手机\"},\n",
    "            {\"price\": {\"$lte\": 1000}}\n",
    "        ]\n",
    "    },\n",
    "    include=[\"embeddings\", \"documents\", \"metadatas\"]  # 移除 \"ids\"\n",
    ")\n",
    "\n",
    "# 检查是否有返回结果，且结果不为 None\n",
    "if not filtered_data or not filtered_data.get('documents'):\n",
    "    print(\"没有找到符合条件的产品。\")\n",
    "else:\n",
    "    # 组合元数据、文档、ID 和嵌入\n",
    "    combined_data = list(zip(\n",
    "        filtered_data.get('metadatas', []),\n",
    "        filtered_data.get('documents', []),\n",
    "        filtered_data.get('ids', []),  # ids 默认返回\n",
    "        filtered_data.get('embeddings', [])\n",
    "    ))\n",
    "\n",
    "    # 定义查询嵌入（例如，用户想查询“高性能手机”）\n",
    "    query_text = \"高性能手机\"\n",
    "    query_embedding = embedding.embed_query(query_text)\n",
    "\n",
    "    # 自定义相似度函数（余弦相似度）\n",
    "    def cosine_similarity(a, b):\n",
    "        return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))\n",
    "\n",
    "    # 对结果根据自定义相似度进行排序\n",
    "    sorted_data = sorted(\n",
    "        combined_data,\n",
    "        key=lambda x: cosine_similarity(query_embedding, x[3]),\n",
    "        reverse=True  # 相似度从高到低排序\n",
    "    )\n",
    "\n",
    "    print(f\"\\n查询 '{query_text}' 的结果（按相似度排序）：\")\n",
    "    for metadata, document, doc_id, _ in sorted_data:\n",
    "        print(f\"ID: {doc_id}, 产品: {document}, 价格: {metadata['price']}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import chromadb\n",
    "from chromadb.utils import embedding_functions\n",
    "from pprint import pprint\n",
    "import os\n",
    "import numpy as np\n",
    "# 假设您已经定义了 ZhipuAIEmbeddings，如果没有，需要替换为实际的嵌入函数\n",
    "# 请确保您已经正确导入或定义了 ZhipuAIEmbeaf77489efe3a6ca0afc5dccdc4d4e657ddings 类\n",
    "\n",
    "# 初始化 Chroma 客户端（持久化模式）\n",
    "client = chromadb.PersistentClient(path=\"ch16_db\")\n",
    "\n",
    "# 删除并重新创建集合\n",
    "try:\n",
    "    client.delete_collection(name=\"products\")\n",
    "    print(\"集合已删除。\")\n",
    "except Exception as e:\n",
    "    print(\"集合不存在，创建新的集合。\")\n",
    "\n",
    "# 创建新的集合\n",
    "collection = client.create_collection(name=\"products\")\n",
    "print(\"新的集合已创建。\")\n",
    "\n",
    "# 定义嵌入函数（使用预训练的嵌入模型）\n",
    "embedding = ZhipuAIEmbeddings(\n",
    "    model=\"embedding-2\",\n",
    "    api_key=os.getenv(\"ZHIPU_API_KEY\"),\n",
    "    dimensions=1024\n",
    ")\n",
    "\n",
    "# 添加嵌入数据\n",
    "documents = [\"Galaxy S21\", \"iPhone 13\", \"MacBook Pro\"]\n",
    "metadatas = [\n",
    "    {\"category\": \"手机\", \"price\": 799.99},\n",
    "    {\"category\": \"手机\", \"price\": 999.99},\n",
    "    {\"category\": \"笔记本电脑\", \"price\": 1299.99}\n",
    "]\n",
    "ids = [\"prod1\", \"prod2\", \"prod3\"]\n",
    "embeddings = embedding.embed_documents(documents)\n",
    "\n",
    "collection.add(\n",
    "    documents=documents,\n",
    "    embeddings=embeddings,\n",
    "    metadatas=metadatas,\n",
    "    ids=ids\n",
    ")\n",
    "print(\"数据添加完成！\")\n",
    "\n",
    "# 打印集合中的所有数据，验证添加是否成功\n",
    "all_data = collection.get(include=[\"documents\", \"metadatas\"])  # ids 会自动返回\n",
    "\n",
    "print(\"当前集合中的数据：\")\n",
    "for metadata, document, doc_id in zip(all_data.get('metadatas', []), all_data.get('documents', []), all_data.get('ids', [])):\n",
    "    print(f\"ID: {doc_id}, 产品: {document}, 分类: {metadata['category']}, 价格: {metadata['price']}\")\n",
    "\n",
    "# 定义自定义相似度检索器\n",
    "class CustomSimilarityRetriever:\n",
    "    def __init__(self, collection, embedding_function, k=3):\n",
    "        self.collection = collection\n",
    "        self.embedding_function = embedding_function\n",
    "        self.k = k\n",
    "\n",
    "    def get_relevant_documents(self, query):\n",
    "        # 计算查询的嵌入\n",
    "        query_embedding = self.embedding_function.embed_query(query)\n",
    "        # 获取集合中的所有嵌入数据\n",
    "        all_data = self.collection.get(\n",
    "            include=[\"embeddings\", \"documents\", \"metadatas\"]  # 移除了 \"ids\"\n",
    "        )\n",
    "        if not all_data or not all_data.get('documents'):\n",
    "            return []\n",
    "        embeddings = all_data.get('embeddings')\n",
    "        documents = all_data.get('documents')\n",
    "        metadatas = all_data.get('metadatas')\n",
    "        ids = all_data.get('ids')  # ids 会自动返回\n",
    "\n",
    "        # 转换嵌入为 numpy 数组\n",
    "        query_embedding = np.array(query_embedding)\n",
    "        doc_embeddings = np.array(embeddings)\n",
    "\n",
    "        # 定义余弦相似度函数\n",
    "        def cosine_similarity(a, b):\n",
    "            return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))\n",
    "\n",
    "        # 计算查询与每个文档的相似度\n",
    "        similarities = [cosine_similarity(query_embedding, doc_embedding) for doc_embedding in doc_embeddings]\n",
    "\n",
    "        # 获取相似度最高的前 k 个文档的索引\n",
    "        top_k_indices = np.argsort(similarities)[-self.k:][::-1]\n",
    "\n",
    "        # 准备返回的文档列表\n",
    "        results = []\n",
    "        for idx in top_k_indices:\n",
    "            doc = {\n",
    "                'document': documents[idx],\n",
    "                'metadata': metadatas[idx],\n",
    "                'id': ids[idx],\n",
    "                # 'embedding': embeddings[idx]  # 如果需要，可以包括嵌入向量\n",
    "            }\n",
    "            results.append(doc)\n",
    "        return results\n",
    "\n",
    "# 创建自定义相似度检索器\n",
    "custom_retriever = CustomSimilarityRetriever(\n",
    "    collection=collection,\n",
    "    embedding_function=embedding,\n",
    "    k=3  # 返回3个文档\n",
    ")\n",
    "\n",
    "# 使用自定义检索器进行检索\n",
    "query_text = \"高性能手机\"\n",
    "retrieved_docs = custom_retriever.get_relevant_documents(query_text)\n",
    "\n",
    "print(f\"\\n查询 '{query_text}' 的结果：\")\n",
    "for doc in retrieved_docs:\n",
    "    print(f\"ID: {doc['id']}, 产品: {doc['document']}, 价格: {doc['metadata']['price']}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from typing import Dict, List\n",
    "import os\n",
    "\n",
    "class SmartPromptGenerator:\n",
    "    def __init__(self, embedding_function):\n",
    "        self.embedding_function = embedding_function\n",
    "        self.categories = {\n",
    "            \"手机\": [\"手机\", \"智能手机\", \"iphone\", \"android\", \"galaxy\"],\n",
    "            \"电脑\": [\"电脑\", \"笔记本\", \"macbook\", \"laptop\", \"notebook\"],\n",
    "            \"平板\": [\"平板\", \"ipad\", \"tablet\"],\n",
    "        }\n",
    "        \n",
    "        # 改进的广告文案模板及其特征描述\n",
    "        self.templates = {\n",
    "            \"手机\": [\n",
    "                {\n",
    "                    \"template\": \"突破创新边界，{product}让科技融入生活\",\n",
    "                    \"features\": \"创新 科技 生活 智能 便携\",\n",
    "                    \"price_range\": \"all\",  # 适用所有价格段\n",
    "                    \"keywords\": [\"创新\", \"智能\", \"科技\"]\n",
    "                },\n",
    "                {\n",
    "                    \"template\": \"非凡性能，精致设计，{product}为您带来智能新体验\",\n",
    "                    \"features\": \"性能 设计 体验 高端 精致\",\n",
    "                    \"price_range\": \"high\",  # 高端价格段\n",
    "                    \"keywords\": [\"性能\", \"设计\", \"高端\"]\n",
    "                },\n",
    "                {\n",
    "                    \"template\": \"震撼视觉，强劲性能，{product}定义科技新高度\",\n",
    "                    \"features\": \"视觉 性能 科技 强劲 高清\",\n",
    "                    \"price_range\": \"high\",\n",
    "                    \"keywords\": [\"视觉\", \"性能\", \"科技\"]\n",
    "                }\n",
    "            ],\n",
    "            \"电脑\": [\n",
    "                {\n",
    "                    \"template\": \"高效办公的得力助手，{product}让工作更轻松\",\n",
    "                    \"features\": \"办公 效率 工作 生产力 便携\",\n",
    "                    \"price_range\": \"medium\",\n",
    "                    \"keywords\": [\"办公\", \"效率\", \"便携\"]\n",
    "                },\n",
    "                {\n",
    "                    \"template\": \"性能之选，{product}为创造力而生\",\n",
    "                    \"features\": \"性能 创造 专业 设计 创意\",\n",
    "                    \"price_range\": \"high\",\n",
    "                    \"keywords\": [\"性能\", \"专业\", \"创意\"]\n",
    "                },\n",
    "                {\n",
    "                    \"template\": \"专业之选，{product}助您突破工作极限\",\n",
    "                    \"features\": \"专业 工作站 性能 效率 极限\",\n",
    "                    \"price_range\": \"high\",\n",
    "                    \"keywords\": [\"专业\", \"性能\", \"极限\"]\n",
    "                }\n",
    "            ],\n",
    "            \"平板\": [\n",
    "                {\n",
    "                    \"template\": \"创意无界，{product}让灵感自由绽放\",\n",
    "                    \"features\": \"创意 灵感 自由 便携 轻薄\",\n",
    "                    \"price_range\": \"all\",\n",
    "                    \"keywords\": [\"创意\", \"灵感\", \"便携\"]\n",
    "                },\n",
    "                {\n",
    "                    \"template\": \"轻盈优雅，{product}让创造更随心\",\n",
    "                    \"features\": \"轻薄 优雅 设计 便携 时尚\",\n",
    "                    \"price_range\": \"high\",\n",
    "                    \"keywords\": [\"轻薄\", \"优雅\", \"时尚\"]\n",
    "                },\n",
    "                {\n",
    "                    \"template\": \"随时随地，{product}让工作娱乐更随心\",\n",
    "                    \"features\": \"便携 娱乐 工作 灵活 多功能\",\n",
    "                    \"price_range\": \"medium\",\n",
    "                    \"keywords\": [\"便携\", \"娱乐\", \"多功能\"]\n",
    "                }\n",
    "            ]\n",
    "        }\n",
    "        \n",
    "        # 预计算所有模板的特征向量\n",
    "        self.template_embeddings = {}\n",
    "        for category, templates in self.templates.items():\n",
    "            self.template_embeddings[category] = []\n",
    "            for template in templates:\n",
    "                embedding = self.embedding_function.embed_query(template[\"features\"])\n",
    "                self.template_embeddings[category].append({\n",
    "                    \"template\": template[\"template\"],\n",
    "                    \"embedding\": embedding,\n",
    "                    \"price_range\": template[\"price_range\"],\n",
    "                    \"keywords\": template[\"keywords\"]\n",
    "                })\n",
    "\n",
    "    def categorize_product(self, product_name: str) -> str:\n",
    "        \"\"\"根据产品名称判断类别\"\"\"\n",
    "        for category, keywords in self.categories.items():\n",
    "            for keyword in keywords:\n",
    "                if keyword.lower() in product_name.lower():\n",
    "                    return category\n",
    "        return \"general\"\n",
    "    \n",
    "    def get_price_range(self, price: float) -> str:\n",
    "        \"\"\"确定价格区间\"\"\"\n",
    "        if price > 1000:\n",
    "            return \"high\"\n",
    "        elif price > 500:\n",
    "            return \"medium\"\n",
    "        return \"low\"\n",
    "\n",
    "    def extract_product_features(self, product_info: dict) -> str:\n",
    "        \"\"\"提取产品特征关键词\"\"\"\n",
    "        category = self.categorize_product(product_info['document'])\n",
    "        features = []\n",
    "        \n",
    "        # 从产品名称中提取特征\n",
    "        product_name = product_info['document'].lower()\n",
    "        if \"pro\" in product_name:\n",
    "            features.extend([\"专业\", \"高端\"])\n",
    "        if \"max\" in product_name:\n",
    "            features.extend([\"大屏\", \"性能\"])\n",
    "        if \"air\" in product_name:\n",
    "            features.extend([\"轻薄\", \"便携\"])\n",
    "            \n",
    "        # 从价格区间提取特征\n",
    "        price = product_info['metadata'].get('price', 0)\n",
    "        price_range = self.get_price_range(price)\n",
    "        if price_range == \"high\":\n",
    "            features.extend([\"高端\", \"奢华\"])\n",
    "        elif price_range == \"medium\":\n",
    "            features.extend([\"中端\", \"性价比\"])\n",
    "        else:\n",
    "            features.extend([\"实惠\", \"经济\"])\n",
    "            \n",
    "        # 从类别添加特征\n",
    "        if category == \"手机\":\n",
    "            features.extend([\"智能\", \"便携\", \"拍照\"])\n",
    "        elif category == \"电脑\":\n",
    "            features.extend([\"效率\", \"性能\", \"办公\"])\n",
    "        elif category == \"平板\":\n",
    "            features.extend([\"轻薄\", \"便携\", \"创意\"])\n",
    "                    \n",
    "        return \" \".join(list(set(features)))  # 去重\n",
    "\n",
    "    def find_best_template(self, product_info: dict) -> str:\n",
    "        \"\"\"使用词向量相似度找到最匹配的广告文案模板\"\"\"\n",
    "        category = self.categorize_product(product_info['document'])\n",
    "        if category == \"general\":\n",
    "            return \"精品之选，{product}品质保证\"\n",
    "        \n",
    "        price = product_info['metadata'].get('price', 0)\n",
    "        price_range = self.get_price_range(price)\n",
    "            \n",
    "        # 提取产品特征并计算嵌入向量\n",
    "        features = self.extract_product_features(product_info)\n",
    "        features_embedding = self.embedding_function.embed_query(features)\n",
    "        \n",
    "        # 筛选符合价格区间的模板\n",
    "        suitable_templates = [\n",
    "            template for template in self.template_embeddings[category]\n",
    "            if template[\"price_range\"] in [\"all\", price_range]\n",
    "        ]\n",
    "        \n",
    "        if not suitable_templates:\n",
    "            return self.templates[category][0][\"template\"]\n",
    "        \n",
    "        # 计算与筛选后模板的相似度\n",
    "        best_similarity = -1\n",
    "        best_template = None\n",
    "        \n",
    "        for template_info in suitable_templates:\n",
    "            # 计算词向量相似度\n",
    "            similarity = self.cosine_similarity(\n",
    "                np.array(features_embedding),\n",
    "                np.array(template_info[\"embedding\"])\n",
    "            )\n",
    "            \n",
    "            # 根据关键词匹配度增加权重\n",
    "            keyword_match = sum(1 for keyword in template_info[\"keywords\"] \n",
    "                              if keyword.lower() in features.lower())\n",
    "            keyword_bonus = keyword_match * 0.1  # 每个匹配关键词增加0.1的权重\n",
    "            \n",
    "            final_similarity = similarity + keyword_bonus\n",
    "            \n",
    "            if final_similarity > best_similarity:\n",
    "                best_similarity = final_similarity\n",
    "                best_template = template_info[\"template\"]\n",
    "            print(f\"最佳相似度: {best_similarity}, 最佳模板: {best_template}\")\n",
    "                \n",
    "        return best_template or self.templates[category][0][\"template\"]\n",
    "        \n",
    "    def cosine_similarity(self, a, b):\n",
    "        \"\"\"计算余弦相似度\"\"\"\n",
    "        return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))\n",
    "\n",
    "    def generate_ad(self, product_info: dict) -> str:\n",
    "        \"\"\"生成广告文案\"\"\"\n",
    "        template = self.find_best_template(product_info)\n",
    "        return template.format(product=product_info['document'])\n",
    "\n",
    "def main():\n",
    "    # 初始化嵌入模型\n",
    "    embedding = ZhipuAIEmbeddings(\n",
    "        model=\"embedding-2\",\n",
    "        api_key=os.getenv(\"ZHIPU_API_KEY\"),\n",
    "        dimensions=1024\n",
    "    )\n",
    "    \n",
    "    # 初始化生成器\n",
    "    generator = SmartPromptGenerator(embedding)\n",
    "    \n",
    "    # 测试数据\n",
    "    test_product = {\n",
    "        'document': 'Iphone 18 Pro max',\n",
    "        'metadata': {'price': 1099.99, 'category': '手机'},\n",
    "        'id': 'prod1'\n",
    "    }\n",
    "    \n",
    "    # 生成广告文案\n",
    "    ad = generator.generate_ad(test_product)\n",
    "    \n",
    "    # 打印结果\n",
    "    print(f\"产品: {test_product['document']}\")\n",
    "    print(f\"提取的特征: {generator.extract_product_features(test_product)}\")\n",
    "    print(f\"生成的广告文案: {ad}\")\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    main()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "chatchat",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
